
import gradio as gr
from PIL import Image


# SDK引入模型
from dwspark.config import Config
from dwspark.models import Text2Picture
# 日志
# 加载系统环境变量：SPARKAI_UID、SPARKAI_APP_ID、SPARKAI_API_KEY、SPARKAI_API_SECRET
config = Config()
model = Text2Picture(config)

img_save_path = 'D://'
def gen_audio(chat_query: str):
    '''
    大模型对话
    :param chat_query: 当前用户问题
    :param chat_history: 历史对话
    :return:
    '''

    img_path = model.generate(chat_query, img_save_path)
    if img_path != '':
        return Image.open(img_path)
    else:
        return None



# gradio页面绘制
with gr.Blocks() as demo:
    # 聊天对话框
    img = gr.Image(label='生成的图片', type='pil')
    # 输入框
    chat_query = gr.Textbox(label="输入问题", placeholder="输入需要咨询的问题")
    # 按钮
    llm_submit_tab = gr.Button("发送", visible=True)
    # 问题样例
    gr.Examples(["画一幅Q版的鲸鱼水墨画"], chat_query)
    # 按钮出发逻辑
    llm_submit_tab.click(fn=gen_audio, inputs=[chat_query], outputs=[img])

if __name__ == "__main__":
    demo.queue().launch()
